'Utilities for working with flv files'

from __future__ import with_statement
import array, struct

def duration(header):
    'Duration of flv file'
    i = header.index('duration') + len('duration')
    a = array.array('d')
    a.fromstring(header[i + 1:i + 9])
    a.byteswap()
    return a[0]

def tag_value(text, tag, offset, length):
    'Value of tag in text'
    i = text.index(tag) + len(tag) + offset
    return text[i:i + length]

def decode_double(s):
    'Double value from binary encoded string'
    a = array.array('d')
    a.fromstring(s)
    a.byteswap()
    return a[0]
    

def replace(text, tag, offset, length, new):
    'Replace value of tag in text with new (which must string)'
    i = text.index(tag) + len(tag)
    return text[:i + offset] + new + text[i + offset + length:]


def header_body(flv_path):
    'Tuple (header, body) of flv file without onMetaData tag'
    with open(flv_path, 'rb') as f:
        header = f.read(9)
        f.read(4)
        f.read(1)
        data_size = struct.unpack('>I', '\x00' + f.read(3))[0]
        f.read(7)
        f.read(data_size)
        body = f.read()
    return header, body

def flv_header(f):
    'First 9 bytes from file'
    return f.read(9)

def skip_metadata(f):
    'File position must be on previous tag size'
    result = ''
    f.read(4)
    result += f.read(1)
    size = f.read(3)
    result += size
    data_size = struct.unpack('>I', '\x00' + size)[0]
    result += f.read(7)
    return result + f.read(data_size)


def flv_metadata(f):
    'File position must be on previous tag size'
    return skip_metadata(f)

def add_prev_size(tag, size):
    'Replace previous tag size with size'
    return struct.pack('>I', size) + tag

def previous_tag_size(f):
    return struct.unpack('>I', f.read(4))


def tag_timestamp(tag):
    'Timestamp value of in milliseconds'
    return struct.unpack('>I', '\x00' + tag[4:7])[0]


def replace_timestamp(tag, timestamp):
    'Replace timestamp in tag with new value'
    packed = struct.pack('>I', timestamp)[1:]
    return tag[:4] + packed + tag[7:]

def simple_merge(out_path, *paths):
    'Merge paths (which must be flv files) and write output to out_path'
    parsed = []
    files = [open(p, 'rb') for p in paths]
    header = flv_header(files[0])
    for f in files[1:]: flv_header(f)
    metadata = [flv_metadata(f) for f in files]
    total = 0
    for meta in metadata:
        total += duration(meta)
    metadata[0] = replace(metadata[0], 'duration', 1, 8, struct.pack('>d', total))
    with open(out_path, 'wb') as out:
        out.write(header)
        prev_size = 0
        offset = 0
        # uncomment next line if you want save metadata
        # out.write(add_prev_size(metadata[0], prev_size))
        for f in files:
            prev = previous_tag_size(f)
            for tag, size in tags(f):
                new_offset = tag_timestamp(tag)
                tag = replace_timestamp(tag, offset + new_offset)
                out.write(add_prev_size(tag, prev_size))
                prev_size = size
                prev = previous_tag_size(f)
            offset += new_offset



def tags(flv, max_tags = -1):
    '''Generator yielding (tag, size) from flv stream. File position must be on
tag type'''
    tag_type = flv.read(1)
    while tag_type and max_tags:
        basic_size = flv.read(3)
        data_size = struct.unpack('>I', '\x00' + basic_size)[0]
        buf = tag_type + basic_size
        buf += flv.read(7)
        buf += flv.read(data_size)
        #print(flv.tell())
        yield buf, data_size + 11
        tag_type = flv.read(1)
        max_tags -= 1

import glob, os.path, re, tkFileDialog, tkMessageBox, webbrowser, sys

def flv_files(directory):
    ''''List of flv files in directory sorted in right order
(i.e part2 is before part10)'''
    lst = glob.glob(os.path.join(directory, '*.flv'))
    def right_cmp(x, y):
        xs = int(re.findall(r'\d+', x)[-1])
        ys = int(re.findall(r'\d+', y)[-1])
        return cmp(xs, ys)
    lst.sort(right_cmp)
    return lst


if __name__ == '__main__':
    directory = tkFileDialog.askdirectory()
    if directory:
        out_file = os.path.join(os.path.dirname(directory),
                                os.path.basename(directory) + '.flv')
        lst = flv_files(directory)
        if len(lst) == 0:
            tkMessageBox.showerror('Error', 'No flv files in %s' % directory)
        else:
            simple_merge(out_file, *lst)
            webbrowser.open(out_file)
